import time
import pysnooper
import pathlib


def count_parameters(model):
    def _model(*args, **kwargs):
        print(f'Total params: {sum(p.numel() for p in model.parameters() if p.requires_grad)}')
        result = model(*args, **kwargs)
        return result
    return _model


def timeit(fn):
    def _fn(*args, **kwargs):
        start = time.time()
        result = fn(*args, **kwargs)
        end = time.time()
        print(f'Time used: {end - start}')
        return result
    return _fn


def all(fn):
    def _fn(*args, **kwargs):
        return timeit(count_parameters(fn))(*args, **kwargs)
    return _fn


def snoop(*args, **kwargs):
    log = '/home/yusongli/Documents/vit-pytorch/debug.log'
    pathlib.Path(log).parent.mkdir(parents=True, exist_ok=True)
    return pysnooper.snoop(
        log,
        *args,
        **kwargs
    )


